package com.ifp.opengate.boot.factory;

import static org.springframework.cloud.gateway.filter.AdaptCachedBodyGlobalFilter.CACHED_REQUEST_BODY_KEY;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.gateway.handler.AsyncPredicate;
import org.springframework.cloud.gateway.handler.predicate.AbstractRoutePredicateFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;

import com.ifp.adapter.common.AdapterCanstants;
import com.ifp.opengate.boot.predicate.ReadBodyPredicate;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * Created by mogs on 2019/4/23.
 */
@Component
public class IFPReadBodyPredicateFactory extends AbstractRoutePredicateFactory<IFPReadBodyPredicateFactory.Config> {
    protected static final Log LOGGER = LogFactory.getLog(IFPReadBodyPredicateFactory.class);

    private static final String TEST_ATTRIBUTE = "read_body_predicate_test_attribute";

    private static final String CACHE_REQUEST_BODY_OBJECT_KEY = "cachedRequestBodyObject";

    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();

    public IFPReadBodyPredicateFactory() {
        super(Config.class);
    }

    @Override
    @SuppressWarnings("unchecked")
    public AsyncPredicate<ServerWebExchange> applyAsync(Config config) {
        return exchange -> {
            final ServerHttpRequest request = exchange.getRequest();
            String methodValue = request.getMethodValue();

            Class inClass = config.getInClass();

            Object cachedBody = exchange.getAttribute(CACHE_REQUEST_BODY_OBJECT_KEY);
            Mono<?> modifiedBody;
            // We can only read the body from the request once, once that happens if we
            // try to read the body again an exception will be thrown. The below if/else
            // caches the body object as a request attribute in the ServerWebExchange
            // so if this filter is run more than once (due to more than one route
            // using it) we do not try to read the request body multiple times
            if (cachedBody != null) {
                try {
                    boolean test = config.getPredicate().test(cachedBody);
                    exchange.getAttributes().put(TEST_ATTRIBUTE, test);
                    return Mono.just(test);
                } catch (ClassCastException e) {
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("Predicate test failed because class in predicate "
                            + "does not match the cached body object", e);
                    }
                }
                return Mono.just(false);
            } else {

                if ("GET".equalsIgnoreCase(methodValue)) {

                    Map<String, Object> requestMap = new HashMap<>();
                    HashMap<Object, Object> head = new HashMap<>();
                    requestMap.put(AdapterCanstants.HEADER, head);
                    requestMap.put("body", request.getQueryParams().toSingleValueMap());
                    exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, requestMap);
                    return Mono.just(true);
                } else {
                    // Join all the DataBuffers so we have a single DataBuffer for the body
                    return DataBufferUtils.join(request.getBody()).flatMap(dataBuffer -> {
                        byte[] bytes = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(bytes);
                        // DataBufferUtils.release(dataBuffer);
                        Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                            DataBufferUtils.retain(buffer);
                            return Mono.just(buffer);
                        });

                        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(request) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedFlux;
                            }
                        };
                        // 解决内存泄露
                        DataBufferUtils.release(dataBuffer);
                        return ServerRequest.create(exchange.mutate().request(mutatedRequest).build(), messageReaders)
                            .bodyToMono(inClass).doOnNext(objectValue -> {
                                exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, objectValue);
                                exchange.getAttributes().put(CACHED_REQUEST_BODY_KEY, cachedFlux);
                            }).map(objectValue -> config.getPredicate().test(objectValue));
                    });
                }

            }
        };
    }

    @Override
    @SuppressWarnings("unchecked")
    public Predicate<ServerWebExchange> apply(Config config) {
        throw new UnsupportedOperationException("ReadBodyPredicateFactory is only async.");
    }

    public static class Config {

        private Class inClass;

        private Predicate predicate = new ReadBodyPredicate();

        private Map<String, Object> hints;

        public Class getInClass() {
            if (null == inClass) {
                inClass = Object.class;
            }
            return inClass;
        }

        public Config setInClass(Class inClass) {
            this.inClass = inClass;
            return this;
        }

        public Predicate getPredicate() {
            if (null == predicate) {
                predicate = new ReadBodyPredicate();
            }
            return predicate;
        }

        public Config setPredicate(Predicate predicate) {
            this.predicate = predicate;
            return this;
        }

        public <T> Config setPredicate(Class<T> inClass, Predicate<T> predicate) {
            setInClass(inClass);
            this.predicate = predicate;
            return this;
        }

        public Map<String, Object> getHints() {
            return hints;
        }

        public Config setHints(Map<String, Object> hints) {
            this.hints = hints;
            return this;
        }

    }
}
